"""
Base routing system implementation
"""
from typing import Dict, List, Optional, Set
from ...core.message import Message

class RoutingRule:
    """基础路由规则类"""
    
    def __init__(self):
        self.routes: Dict[str, str] = {}  # 点对点路由规则 {source: target}
        self.default_order: List[str] = []  # 默认路由顺序
        
    def add_route(self, source: str, target: str) -> None:
        """
        添加一个点对点路由规则
        Args:
            source: 源代理ID
            target: 目标代理ID
        """
        self.routes[source] = target
        
    def set_default_order(self, order: List[str]) -> None:
        """
        设置默认的消息传递顺序
        Args:
            order: 代理ID列表，按处理顺序排列
        """
        self.default_order = order.copy()
        
    def add_to_default_order(self, agent_id: str) -> None:
        """
        将代理添加到默认顺序的末尾
        Args:
            agent_id: 要添加的代理ID
        """
        if agent_id not in self.default_order:
            self.default_order.append(agent_id)
            
    def get_next_agent(self, current_agent: str) -> Optional[str]:
        """
        获取下一个应该处理消息的代理
        Args:
            current_agent: 当前代理ID
        Returns:
            下一个代理ID，如果没有下一个代理则返回None
        """
        # 优先检查特定路由规则
        if current_agent in self.routes:
            return self.routes[current_agent]
            
        # 否则使用默认顺序
        try:
            current_index = self.default_order.index(current_agent)
            # if current_index < len(self.default_order) - 1:
            return self.default_order[(current_index + 1) % len(self.default_order)]
        except ValueError:
            pass
        return None
        
    def get_all_agents(self) -> Set[str]:
        """
        获取所有涉及到的代理ID
        Returns:
            代理ID集合
        """
        agents = set(self.default_order)
        agents.update(self.routes.keys())
        agents.update(self.routes.values())
        return agents
        
    def validate_agents(self, available_agents: Set[str]) -> bool:
        """
        验证路由规则中的所有代理是否都存在
        Args:
            available_agents: 可用的代理ID集合
        Returns:
            如果所有代理都存在则返回True，否则返回False
        """
        for agent_id in self.get_all_agents():
            if agent_id not in available_agents:
                return False
        return True
        
    def clear(self) -> None:
        """清除所有路由规则"""
        self.routes.clear()
        self.default_order.clear()

class MessageRouter:
    """消息路由器，整合各种路由规则"""
    
    def __init__(self):
        self.base_rule = RoutingRule()
        
    def add_agent(self, agent_id: str) -> None:
        """
        添加代理到路由系统
        Args:
            agent_id: 代理ID
        """
        self.base_rule.add_to_default_order(agent_id)
        
    def set_route(self, source: str, target: str) -> None:
        """
        设置点对点路由规则
        Args:
            source: 源代理ID
            target: 目标代理ID
        """
        self.base_rule.add_route(source, target)
        
    def set_default_order(self, order: List[str]) -> None:
        """
        设置默认的消息传递顺序
        Args:
            order: 代理ID列表
        """
        self.base_rule.set_default_order(order)
        
    def get_next_receiver(self, message: Message) -> Optional[str]:
        """
        确定消息的下一个接收者
        Args:
            message: 消息对象
        Returns:
            下一个接收者的ID，如果没有下一个接收者则返回None
        """
        # 如果消息已指定接收者，则返回该接收者
        if message.receiver:
            return message.receiver
            
        # 否则根据路由规则确定下一个接收者
        return self.base_rule.get_next_agent(message.sender)
        
    def validate_routes(self, available_agents: Set[str]) -> bool:
        """
        验证所有路由规则是否有效
        Args:
            available_agents: 可用的代理ID集合
        Returns:
            如果所有路由规则有效则返回True，否则返回False
        """
        return self.base_rule.validate_agents(available_agents)
